与 CUDA 相比,Taichi 跑得快不快?
在「v1.0.0」用 Taichi AOT 方案将特效部署到移动端中,我们展示了 Taichi 脱离 Python 环境部署的方法,实现了将特效算法快速部署在安卓手机上,并加上了重力交互功能。受制于散热、成本等问题,在部署环境中的硬件其计算能力通常很弱,因此有一个问题变得非常重要:相比于更底层的原生编程语言,Taichi 能不能充分地发挥设备的计算能力?于是我们发起 Taichi benchmark 项目做一个更加细分、准确的性能评估。
作为一个领域编程语言(domain-specific language),Taichi 可以用很短的代码解决数值计算方面的问题。我们针对领域中常用算法建立了测试集,将每一个测试项目和性能最好的实现方法进行对比,简而言之:对比基线是外部最强实现 vs Taichi 实现。通过这样的对比,我们试图评估 Taichi 编译器中内置优化技术的有效性以及可优化的空间。同时,这种对比也可以帮助改进 Taichi,以达到更好的性能。作为这个项目的开端,目前版本的性能评测主要是 Taichi 与 CUDA,通过与高度优化的 CUDA 代码进行对照,以评估 Taichi 的优化水平。
这次评测部分地解答了用户一直以来的疑问:Taichi 和 CUDA 比性能究竟怎么样?我们在一块 Nvidia Geforce RTX 3080 上进行了测试,其中有 9 个算法以高性能的 CUDA 实现作为参考。在测试中计时方法对结果很重要,采用的两种计时方式分别是:
外部计时:重复运行计算函数,取平均计算时间。其中第一次运行包含了 Taichi 的编译开销,不计入总时间。
Kernel 计时:在 CUDA 中使用 cudaEventElapsedTime 接口,在 Taichi 中使用 ti.profiler ¹ 统计 kernel 的运行时间。二者的底层实现统一调用 CUPTI,因此在度量上是一致的。
其中,外部计时的方式仍然包含了每次调用的 Python 开销,更贴近实际的使用场景。Kernel 计时仅包含 GPU 计算函数,能更准确地评估 Kernel 函数的质量。无论采用哪种方式,在每个测试项目中我们使用相同的计时器以保证对比是公平的。
下图中我们统计了 Taichi 在不同测试算法中相对 CUDA 的加速比。
在大部分的测试中,Taichi 取得了和 CUDA 非常接近的性能,有些稍好,有些稍差。在 MPM 测试中,Taichi 的性能显著高于 CUDA。这是由于 MPM 算法相对复杂,Taichi 编译器能够自动发现一些优化模式,从而达到比手工优化更好的性能。
另外,我们在 SAXPY 测试中应用了天花板模型(Roofline Model)。天花板模型是高性能计算中很常用的评估芯片计算能力上限的方式,算法的计算访存比越高,就越不容易受内存带宽的限制,设备能发挥出的计算能力就越多。在测试中,我们对标准 SAXPY 算法做了一点修改,在每次内存访问的时候增加多次乘加运算,从而得到了具有不同算术密度的算法。如图所示,Taichi 在各个参数下均达到了非常接近天花板的性能,说明在这种规整的计算下,Taichi 的并行优化能力足够强,可以充分地发挥设备的计算能力。
除了正面的结果,Taichi 在多体模拟测试(图 1 的 N-body)中遇到了比较大的性能问题。在测试中我们对比了一个非常详细的 CUDA 参考实现 ²,它在不同的优化步骤上都给出了对应的 CUDA 代码。Taichi 的性能可以接近优化的中间版本,但无法达到最优的结果。我们发现主要问题是 Taichi 没有支持 CUDA 的 float4,同时对 shared memory 的支持也有欠缺。下一步,Taichi 会持续改进编译器,集成更多的优化技术,赋予程序更好的性能。
更详细的性能分析请参考 Taichi benchmark 代码仓库 ³,阅读原文即可前往。另外在每个子目录里,我们准备了一个独立的 plot_benchmark.py 脚本,只要装好了 nvcc 和 Taichi 就可以一步复现所有的性能图表。如果你有更多想加入 Benchmark 的算法项目,或者性能更强的 baseline,可以在代码仓库中提交 issue。欢迎大家一起来测试 Taichi!
References:1. Taichi profiler 文档 https://docs.taichi.graphics/lang/articles/profiler2. N-body 参考代码仓库 https://github.com/harrism/mini-nbody3. Taichi benchmark 代码仓库 https://github.com/taichi-dev/taichi_benchmark